#!/usr/bin/env python
# -*- coding: utf-8 -*-
import traceback


import numpy as np
import pmt
from gnuradio import gr
from gnuradio import uhd
from gnuradio import blocks
from gnuradio import digital
from gnuradio import filter
import time
import os
import threading
import matplotlib.pyplot as plt
import glob
from scipy import linalg


class BidirectionalChannelEstimation(gr.top_block):
    def __init__(self):
        gr.top_block.__init__(self, "双向信道估计系统", catch_exceptions=True)

        # 硬件配置参数
        self.samp_rate = 3e5  # 采样速率
        self.center_freq = 915e6  # 中心频率
        self._tx_gain = 40  # 增益
        self._rx_gain = 60  # 增益
        self.sps = 4  # 每符号采样数
        self.pilot_length = 127  # 导频序列长度
        self.num_repeated_pilot = 5  # 重复发送导频的次数
        self.num_iterations = 10  # 探测次数
        self.interval = 0.2  # 信道探测间隔时间
        self.current_role = "tx_rx"  # 当前角色
        self._running = False  # 运行状态标志
        self._tx_delay = 0.06  # 发射时间
        self._rx_delay = 0.01  # 接收时间
        self.switch_tx_rx_duration = 0.001
        self._tx_buffer_size = 512
        self._rx_buffer_size = 512
        # CSI数据存储
        self.csi_results = {
            'dir1': {'time': [], 'freq': [], 'raw': [],'sync_offset':[]},  # 方向1的CSI数据
            'dir2': {'time': [], 'freq': [], 'raw': [],'sync_offset':[]}  # 方向2的CSI数据
        }
        # 生成Gold序列
        # self.gold_bits = self.generate_gold_sequence(6, [6, 1], [6, 5, 2, 1])
        # self.gold_code = ''.join(map(str, self.gold_bits))

        # 生成pilot序列

        self.pilot_bits = np.random.randint(0,2,size=(self.pilot_length,))
        # 将PN序列转换为复数符号(BPSK调制)
        self.ref_symbols = np.array([1.0 if bit else -1.0 for bit in self.pilot_bits],
                                    dtype=np.complex64)
        # self.transmitting_symbols = np.concatenate(
        #     [self.gold_bits, np.tile(self.pilot_bits, self.num_repeated_pilot)])  # 重复参考信号多次
        self.transmitting_symbols = np.tile(self.pilot_bits, self.num_repeated_pilot)  # 重复参考信号多次

        # 初始化流图组件
        self.pilot_src = blocks.vector_source_b(self.transmitting_symbols, False)
        self.mod_bpsk = digital.chunks_to_symbols_bc([1.0 + 0j, -1.0 + 0j])  # 符号映射

        # 根升余弦滤波器
        self.rrc_tx_filter = filter.interp_fir_filter_ccf(
            self.sps,
            filter.firdes.root_raised_cosine(1.0, self.sps, 1.0, 0.35, 41)
        )
        self.rrc_rx_filter = filter.fir_filter_ccf(
            1,  # 接收端不需要插值
            filter.firdes.root_raised_cosine(1.0, self.sps, 1.0, 0.35, 41))

        # 完成下采样和最佳采样点选择
        self.timing_recovery = digital.symbol_sync_cc(
            digital.TED_GARDNER,  # 使用Gardner定时误差检测器
            sps=self.sps,  # 每符号采样数（与发送端一致）
            loop_bw=0.05,  # 环路带宽（控制收敛速度）
            damping_factor=1.0,  # 阻尼系数（影响稳定性）
            ted_gain=1.0,  # TED增益
            max_deviation=1.5,  # 最大时钟偏差容限
            osps=1,
            slicer=digital.constellation_bpsk().base()
        )

        # BPSK解调器
        self.demod_BPSK = digital.constellation_decoder_cb(digital.constellation_bpsk().base())
        self.rx_sink = blocks.vector_sink_c()  # 接收数据存储
        self.vec_sink = blocks.vector_sink_b()
        self.debug_sink = blocks.vector_sink_f()

        # 同步检测
        # self.sync_corr = digital.correlate_access_code_tag_bb(self.gold_code, 17, "sync_pos")
        # self.tag_debug = custom_tag_debug(name="sync_pos")


        # 初始化USRP设备(快速切换优化)
        self.init_usrp_fast_switch()

        # 创建结果保存目录
        os.makedirs("csi_results", exist_ok=True)

    def init_usrp_fast_switch(self):
        """初始化USRP设备(优化快速切换)"""
        # 设备1 (TX/RX)
        self.usrp1_tx = uhd.usrp_sink(
            device_addr="serial=31A92E2",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._tx_buffer_size}",
            )
        )
        self.usrp1_tx.set_antenna("TX/RX", 0)
        self.usrp1_tx.set_gain(self._tx_gain)

        self.usrp1_rx = uhd.usrp_source(
            device_addr="serial=31A92E2",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"recv_frame_size={self._rx_buffer_size}",
            )
        )
        self.usrp1_rx.set_antenna("RX2", 0)
        self.usrp1_rx.set_gain(self._rx_gain)

        # 设备2 (TX/RX)
        self.usrp2_tx = uhd.usrp_sink(
            device_addr="serial=31AC8A5",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"send_frame_size={self._tx_buffer_size}",
            )
        )
        self.usrp2_tx.set_antenna("TX/RX", 0)
        self.usrp2_tx.set_gain(self._tx_gain)

        self.usrp2_rx = uhd.usrp_source(
            device_addr="serial=31AC8A5",
            stream_args=uhd.stream_args(
                cpu_format="fc32",
                args=f"recv_frame_size={self._rx_buffer_size}",
            )
        )
        self.usrp2_rx.set_antenna("RX2", 0)
        self.usrp2_rx.set_gain(self._rx_gain)

        # 配置通用参数(快速调谐优化)
        for usrp in [self.usrp1_tx, self.usrp1_rx, self.usrp2_tx, self.usrp2_rx]:
            usrp.set_samp_rate(self.samp_rate)
            usrp.set_center_freq(uhd.tune_request(self.center_freq))
            usrp.set_clock_source("internal")
            usrp.set_time_source("internal")
            usrp.set_time_now(uhd.time_spec(0.0))
            usrp.set_command_time(uhd.time_spec(0.0))  # 启用定时命令
        time.sleep(0.1)

    def setup_flowgraph(self):
        """超快速流图配置(<2ms)"""
        self.lock()
        try:
            self.disconnect_all()

            if self.current_role == "tx_rx":
                # USRP1作为发射机，USRP2作为接收机
                self.connect(
                    (self.pilot_src, 0),
                    (self.mod_bpsk, 0),
                    (self.rrc_tx_filter, 0),
                    (self.usrp1_tx, 0),
                )
                self.connect(
                    (self.usrp2_rx, 0),
                    (self.rrc_rx_filter, 0),
                    self.timing_recovery,
                    (self.rx_sink, 0)
                )
            else:
                # USRP2作为发射机，USRP1作为接收机
                self.connect(
                    (self.pilot_src, 0),
                    (self.mod_bpsk, 0),
                    (self.rrc_tx_filter, 0),
                    (self.usrp2_tx, 0),
                )
                self.connect(
                    (self.usrp1_rx, 0),
                    (self.rrc_rx_filter, 0),
                    self.timing_recovery,
                    (self.rx_sink, 0)
                )

            # 调试连接
            # self.connect(
            #     (self.timing_recovery, 0),
            #     (self.demod_BPSK, 0),
            #     (self.sync_corr, 0),
            #     (self.tag_debug, 0)
            # )

            self.rx_sink.reset()

        finally:
            self.unlock()

    def switch_roles(self):
        """立即切换收发角色（无延迟）"""
        if self.current_role == "tx_rx":
            # USRP1立即停止发射，USRP2立即停止接收
            self.usrp1_tx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.usrp2_rx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.current_role = "rx_tx"
        else:
            # USRP2立即停止发射，USRP1立即停止接收
            self.usrp2_tx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.usrp1_rx.set_stream_args(uhd.stream_args("fc32", args="flush=1"))
            self.current_role = "tx_rx"

        # 立即重建流图连接
        self.setup_flowgraph()
        time.sleep(self.switch_tx_rx_duration)

    def estimate_channel(self, direction):
        # 设置精确的发射时间
        if direction == "dir1":
            tx_time = self.usrp1_tx.get_time_now() + uhd.time_spec(self._tx_delay)
            rx_time = self.usrp2_rx.get_time_now() + uhd.time_spec(self._rx_delay)
            self.usrp1_tx.set_start_time(tx_time)
            self.usrp2_rx.set_start_time(rx_time)  # 接收同步!
        else:
            tx_time = self.usrp2_tx.get_time_now() + uhd.time_spec(self._tx_delay)
            rx_time = self.usrp1_rx.get_time_now() + uhd.time_spec(self._rx_delay)
            self.usrp2_tx.set_start_time(tx_time)
            self.usrp1_rx.set_start_time(rx_time)  # 接收同步!

        # 精确计算需要的采集时间
        expected_duration = (len(self.transmitting_symbols) * self.sps / self.samp_rate) * 1.2

        # 开始采集
        self.rx_sink.reset()
        # self.tag_debug.clear_tags()
        self.start()

        # 优化的采集窗口
        time.sleep(self._tx_delay + expected_duration)
        self.stop()
        self.wait()

        # 接收数据并分段处理
        rx_data = np.array(self.rx_sink.data())
        if len(rx_data) < len(self.transmitting_symbols):  # 检查是否收到完整数据
            print("数据不完整")
            return None, None

        # print(f'transmitting symbol number: {len(self.transmitting_symbols)}')
        # print(f'receiving symbol number: {len(rx_data)}')
        # print(f'maximum receving signal power: {20 * np.log10(np.max(np.abs(rx_data))) + 30:.2f}dBm')
        # print(f'mean receving signal power: {20 * np.log10(np.mean(np.abs(rx_data))) + 30:.2f}dBm')

        # sync_point = self.get_sync_info()
        # if sync_point is None:
        #     print("未检测到同步码！")
        #     return None, None
        #
        # sync_index = sync_point['sample_index']
        # print('sync_index:', sync_index)

        sync_index = int(self.samp_rate/self.sps*(self._tx_delay-self._rx_delay))
        rx_data_align = rx_data[sync_index:sync_index + self.num_repeated_pilot * self.pilot_length]

        if len(rx_data_align)<self.num_repeated_pilot*self.pilot_length:
            print("同步位置错误！")
            return
        # 将数据分割为 self.num_repeated_pilot 段
        segments = np.reshape(rx_data_align, (self.num_repeated_pilot, self.pilot_length))
        # 时域相干累加
        y = np.mean(segments, axis=0)  # 对多次重复取平均

        # LS估计
        h_est = np.vdot(y, self.ref_symbols) / np.linalg.norm(self.ref_symbols) ** 2
        # csi = np.fft.fft(h_est, n=1024)
        csi = 0

        # 存储结果
        self.csi_results[direction]['time'].append(h_est)
        self.csi_results[direction]['freq'].append(csi)
        self.csi_results[direction]['raw'].append(rx_data)
        self.csi_results[direction]['sync_offset'].append(sync_index)

    def save_results(self):
        """保存所有缓冲的结果"""
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        for direction in ['dir1', 'dir2']:
            if len(self.csi_results[direction]['time']) > 0:
                filename = f"csi_results/{direction}_csi_{timestamp}.npz"
                np.savez(
                    filename,
                    time_response=np.array(self.csi_results[direction]['time']),
                    freq_response=np.array(self.csi_results[direction]['freq']),
                    raw=np.array(self.csi_results[direction]['raw'][0]),
                    sync_offset=np.array(self.csi_results[direction]['sync_offset'][0]),
                )
                print(f"已保存{len(self.csi_results[direction]['time'])}个{direction}结果")

    def run(self):
        """主运行循环"""
        self._running = True
        try:
            # 初始连接配置
            self.setup_flowgraph()

            for i in range(self.num_iterations):
                if not self._running:
                    break

                print(f"\n=== 第 {i + 1}/{self.num_iterations} 次迭代 ===")

                # 方向1测量
                print("[USRP1→USRP2] 开始测量...")
                self.estimate_channel("dir1")

                # 切换方向
                self.switch_roles()

                # 方向2测量
                print("[USRP2→USRP1] 开始测量...")
                self.estimate_channel("dir2")

                # 恢复初始状态
                self.switch_roles()

                if i < self.num_iterations - 1:
                    time.sleep(self.interval)

            # 保存所有结果
            self.save_results()

        except KeyboardInterrupt:
            print("用户中断-保存已收集数据...")
            self.save_results()
        except Exception as e:
            print(f"运行时错误: {str(e)} - 保存已收集数据...")
            traceback.print_exc()
            self.save_results()
        finally:
            self._running = False
            print("测量未完成")

    def generate_gold_sequence(self, n, taps1, taps2, seed=1):
        """
        生成 Gold 序列

        参数：
            n: 移位寄存器位数 (生成序列长度 = 2^n - 1)
            taps1: 第一个 m 序列的反馈抽头 (从1开始计数，如[5,2]表示x^5 + x^2 + 1)
            taps2: 第二个 m 序列的反馈抽头
            seed: 初始寄存器值

        返回：
            numpy 数组 (0/1序列)
        """
        length = 2 ** n - 1

        def generate_mseq(taps):
            reg = seed
            seq = []
            for _ in range(length):
                feedback = 0
                for tap in taps:
                    feedback ^= (reg >> (n - tap)) & 1
                seq.append(reg & 1)
                reg = ((reg << 1) | feedback) & ((1 << n) - 1)
            return np.array(seq)

        # 生成两个 m 序列
        mseq1 = generate_mseq(taps1)
        mseq2 = generate_mseq(taps2)

        # 通过异或生成 Gold 序列
        gold_seq = np.logical_xor(mseq1, mseq2).astype(int)
        return gold_seq

    def get_sync_info(self):
        """获取同步位置和相关性"""
        sync_points = []
        for tag in self.tag_debug.current_tags():
            if pmt.symbol_to_string(tag.key) == "sync_pos":  # 确保标签键匹配
                sync_points.append({
                    'sample_index': tag.offset,
                    'correlation': pmt.to_float(tag.value),
                    'timestamp': tag.offset / self.samp_rate
                })
        return sync_points[0] if sync_points else None  # 返回第一个同步点或None



from gnuradio import gr
import pmt
import numpy as np

class custom_tag_debug(gr.sync_block):
    def __init__(self, name="tag_debug"):
        gr.sync_block.__init__(
            self,
            name="Custom Tag Debug",
            in_sig=[np.byte],  # 使用numpy数据类型指定输入为字节
            out_sig=None       # 无输出
        )
        self.tag_name = name
        self.stored_tags = []
        self.set_tag_propagation_policy(gr.TPP_DONT)

    def work(self, input_items, output_items):
        tags = self.get_tags_in_window(0, 0, len(input_items[0]))
        # print(f"[DEBUG] 收到标签数量: {len(tags)}")  # 打印标签数量
        for tag in tags:
            # print(f"[DEBUG] 标签: key={pmt.symbol_to_string(tag.key)}, value={pmt.to_python(tag.value)}")
            if pmt.symbol_to_string(tag.key) == self.tag_name:
                self.stored_tags.append(tag)
        return len(input_items[0])

    def current_tags(self):
        return self.stored_tags

    def clear_tags(self):
        self.stored_tags = []

def find_latest_csi_files():
    """查找并返回最新的匹配CSI文件对路径

    返回:
        tuple: 返回最新的(dir1文件路径, dir2文件路径)，如果没有匹配对则返回(None, None)
    """
    try:
        # 获取所有CSI文件并按修改时间排序(最新的在前面)
        dir1_files = sorted(glob.glob(os.path.join("csi_results", "dir1_csi_*.npz")),
                            key=os.path.getmtime, reverse=True)
        dir2_files = sorted(glob.glob(os.path.join("csi_results", "dir2_csi_*.npz")),
                            key=os.path.getmtime, reverse=True)

        # 如果没有文件，直接返回None
        if not dir1_files or not dir2_files:
            print("未找到足够的CSI文件")
            return None, None

        # 创建一个字典来存储时间戳和文件路径的映射
        timestamp_map = {}

        # 处理dir1文件
        for file_path in dir1_files:
            try:
                # 从文件名提取时间戳部分，例如：dir1_csi_20230815_143022.npz
                filename = os.path.basename(file_path)
                timestamp = '_'.join(filename.split('_')[2:]).replace('.npz', '')
                timestamp_map[timestamp] = {'dir1': file_path}
            except Exception as e:
                print(f"处理文件{file_path}时出错: {str(e)}")
                continue

        # 处理dir2文件并尝试匹配
        for file_path in dir2_files:
            try:
                filename = os.path.basename(file_path)
                timestamp = '_'.join(filename.split('_')[2:]).replace('.npz', '')

                if timestamp in timestamp_map:
                    timestamp_map[timestamp]['dir2'] = file_path
                    # 找到第一个匹配的就返回(因为文件已经按时间排序)
                    return (timestamp_map[timestamp]['dir1'],
                            timestamp_map[timestamp]['dir2'])
            except Exception as e:
                print(f"处理文件{file_path}时出错: {str(e)}")
                continue

        print("未找到时间戳匹配的CSI文件对")
        return None, None

    except Exception as e:
        print(f"查找CSI文件时发生错误: {str(e)}")
        return None, None


def plot_csi_comparison(dir1_file, dir2_file):
    """绘制两个方向的CSI幅度对比图

    参数:
        dir1_file (str): 方向1的CSI文件路径
        dir2_file (str): 方向2的CSI文件路径
    """

    if dir1_file is None or dir2_file is None:
        print("错误: 必须提供两个有效的文件路径")
        return

    # 加载CSI数据
    dir1_data = np.load(dir1_file)
    dir2_data = np.load(dir2_file)

    # 获取时域信道
    dir1_time = dir1_data['time_response']
    dir2_time = dir2_data['time_response']

    # 计算幅度(dB)
    dir1_mag = 20 * np.log10(1e-16 + np.abs(dir1_time.reshape((-1,))))
    dir2_mag = 20 * np.log10(1e-16 + np.abs(dir2_time.reshape((-1,))))



    # 绘制图形
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
    plt.figure(figsize=(12, 8))

    # 绘制两个方向的CSI幅度对比
    plt.plot(dir1_mag, label='USRP1→USRP2')
    plt.plot(dir2_mag, label='USRP2→USRP1', alpha=0.7)
    plt.title('CSI幅度对比')
    plt.xlabel('时间步')
    plt.ylabel('幅度 (dB)')
    plt.grid(True)
    plt.legend()

    # 相关系数计算
    if len(dir1_mag) < 2 or len(dir2_mag) < 2:
        print("数据长度不足，无法计算相关系数")
        correlation = 0
    else:
        with np.errstate(divide='ignore', invalid='ignore'):
            correlation = np.corrcoef(dir1_mag, dir2_mag)[0, 1]
            if np.isnan(correlation):
                print('correlation is NaN!')
                correlation = 0

    plt.text(0.05, 0.95, f'相似度(相关系数): {correlation:.4f}',
             transform=plt.gca().transAxes, bbox=dict(facecolor='white', alpha=0.8))

    # 保存图形
    timestamp = time.strftime("%Y%m%d_%H%M%S")
    plot_filename = f"csi_results/csi_comparison_{timestamp}.png"
    plt.savefig(plot_filename)
    plt.show()

    print(f"CSI对比图已保存为: {plot_filename}")

    # 打印相似度结论
    if correlation > 0.8:
        print("结论: 两个方向的CSI幅度高度相似")
    elif correlation > 0.5:
        print("结论: 两个方向的CSI幅度中等相似")
    else:
        print("结论: 两个方向的CSI幅度差异较大")

    dir1_raw = dir1_data['raw']
    sync_offset = dir1_data['sync_offset']  # 确保这个值已正确保存和加载

    plt.figure(figsize=(12, 6))
    plt.plot(np.real(dir1_raw), label='接收符号(实部)')

    # 添加同步位置的红色垂直线
    plt.axvline(x=sync_offset, color='r', linestyle='--', linewidth=2,
                label=f'开始位置 (符号: {sync_offset})')
    plt.axvline(x=sync_offset+127*5, color='r', linestyle='--', linewidth=2,
                label=f'结束位置 (符号: {sync_offset+127*5})')

    plt.title('接收符号波形与同步位置')
    plt.xlabel('符号索引')
    plt.ylabel('符号幅度')
    plt.grid(True)
    plt.legend()
    plt.show()


def main():
    print("=== 快速切换双向信道估计系统 ===")
    print("初始化硬件...")

    try:
        estimator = BidirectionalChannelEstimation()
        print("系统准备就绪，开始测量...")
        estimator.run()
    except Exception as e:
        print(f"系统初始化失败: {str(e)}")
        traceback.print_exc()


if __name__ == "__main__":
    main()

    # 获取最新的文件对
    dir1_path, dir2_path = find_latest_csi_files()
    # 绘制对比图
    if dir1_path and dir2_path:
        plot_csi_comparison(dir1_path, dir2_path)
    else:
        print("无法找到有效的CSI文件对")